// middleware/jwt/jwt.go
package middleware

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
	"github.com/golang-jwt/jwt/v5"
	"github.com/spf13/viper"
)

// Claims 定义JWT的声明结构
type Claims struct {
	UserID string `json:"user_id"`
	Role   string `json:"role"`
	jwt.RegisteredClaims
}

// TokenPair 定义令牌对结构
type TokenPair struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
}

// IJWTMiddleware 定义JWT中间件接口
type IJWTMiddleware interface {
	GenerateTokenPair(userID string, role string) (*TokenPair, error)
	AuthMiddleware() gin.HandlerFunc
	RefreshToken(refreshToken string, role string) (*TokenPair, error)
	InvalidateTokens(userID string) error
}

// jwtMiddleware 结构体定义
type jwtMiddleware struct {
	secretKey          []byte
	accessTokenExpire  time.Duration
	refreshTokenExpire time.Duration
	issuer             string
	redis              *redis.Client
}

// NewJWTMiddleware 创建新的JWT中间件实例
func NewJWTMiddleware(redis *redis.Client) (IJWTMiddleware, error) {
	secretKey := viper.GetString("jwt.secret_key")
	if secretKey == "" {
		return nil, errors.New("JWT secret key is not set")
	}

	return &jwtMiddleware{
		secretKey:          []byte(secretKey),
		accessTokenExpire:  viper.GetDuration("jwt.access_token_expire"),
		refreshTokenExpire: viper.GetDuration("jwt.refresh_token_expire"),
		issuer:             viper.GetString("jwt.issuer"),
		redis:              redis,
	}, nil
}

// generateAccessToken 生成访问令牌
func (j *jwtMiddleware) generateAccessToken(userID string, role string) (string, error) {
	now := time.Now()
	claims := Claims{
		UserID: userID,
		Role:   role,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(now.Add(j.accessTokenExpire)),
			IssuedAt:  jwt.NewNumericDate(now),
			Issuer:    j.issuer,
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(j.secretKey)
}

// generateRefreshToken 生成刷新令牌
func (j *jwtMiddleware) generateRefreshToken(userID string) (string, error) {
	now := time.Now()
	claims := Claims{
		UserID: userID,
		RegisteredClaims: jwt.RegisteredClaims{
			ExpiresAt: jwt.NewNumericDate(now.Add(j.refreshTokenExpire)),
			IssuedAt:  jwt.NewNumericDate(now),
			Issuer:    j.issuer,
		},
	}

	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	return token.SignedString(j.secretKey)
}

// GenerateTokenPair 生成访问令牌和刷新令牌对
func (j *jwtMiddleware) GenerateTokenPair(userID string, role string) (*TokenPair, error) {
	// 生成访问令牌
	accessToken, err := j.generateAccessToken(userID, role)
	if err != nil {
		return nil, fmt.Errorf("generate access token error: %w", err)
	}

	// 生成刷新令牌
	refreshToken, err := j.generateRefreshToken(userID)
	if err != nil {
		return nil, fmt.Errorf("generate refresh token error: %w", err)
	}

	// 将刷新令牌存储在Redis中
	err = j.redis.Set(context.Background(),
		fmt.Sprintf("refresh_token:%d", userID),
		refreshToken,
		j.refreshTokenExpire).Err()
	if err != nil {
		return nil, fmt.Errorf("store refresh token error: %w", err)
	}

	return &TokenPair{
		AccessToken:  accessToken,
		RefreshToken: refreshToken,
	}, nil
}

// AuthMiddleware 认证中间件
func (j *jwtMiddleware) AuthMiddleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		token := c.GetHeader("Authorization")
		if token == "" {
			c.JSON(401, gin.H{"error": "no token provided"})
			c.Abort()
			return
		}

		if len(token) > 7 && token[:7] == "Bearer " {
			token = token[7:]
		}

		// 解析令牌
		claims := &Claims{}
		parsedToken, err := jwt.ParseWithClaims(token, claims, func(token *jwt.Token) (interface{}, error) {
			return j.secretKey, nil
		})

		if err != nil {
			if err.Error() == "Token is expired" {
				c.JSON(401, gin.H{"error": "token out of date", "need_refresh": true})
			} else {
				c.JSON(401, gin.H{"error": "invalid token"})
			}
			c.Abort()
			return
		}

		if !parsedToken.Valid {
			c.JSON(401, gin.H{"error": "invalid token"})
			c.Abort()
			return
		}

		// 将用户信息存储在上下文中
		c.Set("user_id", claims.UserID)
		c.Set("role", claims.Role)
		c.Next()
	}
}

// RefreshToken 过期后申请新令牌
func (j *jwtMiddleware) RefreshToken(refreshToken string, role string) (*TokenPair, error) {
	claims := &Claims{}
	parsedToken, err := jwt.ParseWithClaims(refreshToken, claims, func(token *jwt.Token) (interface{}, error) {
		return j.secretKey, nil
	})

	if err != nil || !parsedToken.Valid {
		return nil, errors.New("invalid refresh token")
	}

	// 验证刷新令牌是否在Redis中
	key := fmt.Sprintf("refresh_token:%s", claims.UserID)
	storedToken, err := j.redis.Get(context.Background(), key).Result()
	if err != nil || storedToken != refreshToken {
		return nil, errors.New("refresh token not found or invalid")
	}

	// 先删除旧的刷新令牌
	err = j.redis.Del(context.Background(), key).Err()
	if err != nil {
		return nil, fmt.Errorf("failed to invalidate old refresh token: %w", err)
	}
	// 生成新的令牌对
	return j.GenerateTokenPair(claims.UserID, role)
}

// InvalidateTokens 使令牌失效
func (j *jwtMiddleware) InvalidateTokens(userID string) error {
	// 从Redis中删除刷新令牌
	err := j.redis.Del(context.Background(),
		fmt.Sprintf("refresh_token:%d", userID)).Err()
	if err != nil {
		return fmt.Errorf("invalidate tokens error: %w", err)
	}
	return nil
}
